#include "compi/context.h"
#include "compi/tag.h"

#include <gtest/gtest.h>

#include <mpi.h>
#include <set>
#include <vector>

using namespace compi;

class TagAllocatorTest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        // 确保MPI已初始化
        int initialized = 0;
        MPI_Initialized(&initialized);
        if (!initialized)
        {
            int provided;
            MPI_Init_thread(nullptr, nullptr, MPI_THREAD_MULTIPLE, &provided);
        }
    }

    void TearDown() override
    {
        // 通过 Context 释放与 communicator 关联的资源
        Context::destroy_for_comm(MPI_COMM_WORLD);
    }
};

/**
 * 测试基本的分配器创建和命名空间管理
 */
TEST_F(TagAllocatorTest, BasicAllocatorAndNamespace)
{
    auto & allocator = Context::for_comm(MPI_COMM_WORLD).tag_allocator();

    // 测试标签上界
    EXPECT_GT(allocator.tag_ub(), 0);

    // 创建命名空间
    auto ns1 = allocator.create_namespace("test_ns1");
    auto ns2 = allocator.create_namespace("test_ns2");
    auto ns3 = allocator.create_namespace(); // 无名命名空间

    EXPECT_NE(ns1.id(), ns2.id());
    EXPECT_NE(ns1.id(), ns3.id());
    EXPECT_NE(ns2.id(), ns3.id());

    EXPECT_EQ(ns1.name(), "test_ns1");
    EXPECT_EQ(ns2.name(), "test_ns2");
    EXPECT_EQ(ns3.name(), "");

    // 测试按名称查找
    auto found_ns1 = allocator.find_namespace("test_ns1");
    EXPECT_EQ(found_ns1.id(), ns1.id());

    // 测试重复名称
    EXPECT_THROW(allocator.create_namespace("test_ns1"), std::runtime_error);

    // 测试查找不存在的命名空间
    EXPECT_THROW(allocator.find_namespace("nonexistent"), std::runtime_error);
}

/**
 * 测试标签分配和释放
 */
TEST_F(TagAllocatorTest, TagAllocationAndRelease)
{
    auto & allocator = Context::for_comm(MPI_COMM_WORLD).tag_allocator();
    auto ns = allocator.create_namespace("alloc_test");

    // 分配一些标签
    std::vector<MessageTag> tags;
    for (int i = 0; i < 10; ++i)
    {
        auto tag = ns.request();
        EXPECT_GT(tag.value, 0);
        EXPECT_LT(tag.value, allocator.tag_ub());
        EXPECT_EQ(tag.ns_id, ns.id());
        EXPECT_TRUE(ns.owns(tag.value));
        tags.push_back(tag);
    }

    // 检查标签唯一性
    std::set<int> tag_values;
    for (const auto & tag : tags)
    {
        EXPECT_TRUE(tag_values.insert(tag.value).second) << "Duplicate tag value: " << tag.value;
    }

    // 释放标签
    for (const auto & tag : tags)
    {
        ns.release(tag);
    }

    // 重新分配应该重用释放的标签
    auto new_tag = ns.request();
    EXPECT_TRUE(tag_values.count(new_tag.value) > 0) << "Tag not reused: " << new_tag.value;
}

/**
 * 测试命名空间隔离
 */
TEST_F(TagAllocatorTest, NamespaceIsolation)
{
    auto & allocator = Context::for_comm(MPI_COMM_WORLD).tag_allocator();
    auto ns1 = allocator.create_namespace("isolation_test1");
    auto ns2 = allocator.create_namespace("isolation_test2");

    // 从两个命名空间分配标签
    std::vector<MessageTag> tags1, tags2;
    for (int i = 0; i < 20; ++i)
    {
        tags1.push_back(ns1.request());
        tags2.push_back(ns2.request());
    }

    // 检查命名空间隔离
    for (const auto & tag1 : tags1)
    {
        EXPECT_TRUE(ns1.owns(tag1.value));
        EXPECT_FALSE(ns2.owns(tag1.value));
        EXPECT_TRUE(allocator.belongs(ns1.id(), tag1.value));
        EXPECT_FALSE(allocator.belongs(ns2.id(), tag1.value));
    }

    for (const auto & tag2 : tags2)
    {
        EXPECT_TRUE(ns2.owns(tag2.value));
        EXPECT_FALSE(ns1.owns(tag2.value));
        EXPECT_TRUE(allocator.belongs(ns2.id(), tag2.value));
        EXPECT_FALSE(allocator.belongs(ns1.id(), tag2.value));
    }

    // 检查标签范围不重叠
    std::set<int> all_tag_values;
    for (const auto & tag1 : tags1)
    {
        EXPECT_TRUE(all_tag_values.insert(tag1.value).second) << "Tag collision: " << tag1.value;
    }
    for (const auto & tag2 : tags2)
    {
        EXPECT_TRUE(all_tag_values.insert(tag2.value).second) << "Tag collision: " << tag2.value;
    }
}

/**
 * 测试引用计数
 */
TEST_F(TagAllocatorTest, ReferenceCountingTest)
{
    auto & allocator = Context::for_comm(MPI_COMM_WORLD).tag_allocator();
    auto ns = allocator.create_namespace("refcount_test");

    auto tag = ns.request();
    int original_value = tag.value;

    // 增加引用计数
    ns.retain(tag);
    ns.retain(tag);

    // 释放一次，标签应该仍然被占用
    ns.release(tag);

    // 分配新标签，应该得到不同的值
    auto new_tag = ns.request();
    EXPECT_NE(new_tag.value, original_value);

    // 继续释放原标签
    ns.release(tag);
    ns.release(tag);

    // 现在原标签应该可以重用
    auto reused_tag = ns.request();
    EXPECT_EQ(reused_tag.value, original_value);
}

/**
 * 测试动态范围扩展
 */
TEST_F(TagAllocatorTest, DynamicRangeExpansion)
{
    auto & allocator = Context::for_comm(MPI_COMM_WORLD).tag_allocator();
    auto ns = allocator.create_namespace("expansion_test");

    // 获取初始统计信息
    auto initial_stats = ns.get_stats();
    EXPECT_GT(initial_stats.total_capacity, 0);
    EXPECT_EQ(initial_stats.range_count, 1);

    // 分配超过初始容量的标签
    std::vector<MessageTag> tags;
    int allocation_count = initial_stats.total_capacity + 50;

    for (int i = 0; i < allocation_count; ++i)
    {
        auto tag = ns.request();
        tags.push_back(tag);
    }

    // 检查范围是否扩展
    auto final_stats = ns.get_stats();
    EXPECT_GT(final_stats.total_capacity, initial_stats.total_capacity);
    EXPECT_GE(final_stats.range_count, initial_stats.range_count);
    EXPECT_EQ(final_stats.allocated_count, allocation_count);

    // 验证所有标签都有效
    for (const auto & tag : tags)
    {
        EXPECT_TRUE(ns.owns(tag.value));
        EXPECT_GT(tag.value, 0);
        EXPECT_LT(tag.value, allocator.tag_ub());
    }
}

/**
 * 测试全局统计信息
 */
TEST_F(TagAllocatorTest, GlobalStatistics)
{
    auto & allocator = Context::for_comm(MPI_COMM_WORLD).tag_allocator();

    auto initial_stats = allocator.get_global_stats();
    EXPECT_EQ(initial_stats.total_namespaces, 0);
    EXPECT_EQ(initial_stats.total_allocated_tags, 0);

    // 创建命名空间并分配标签
    auto ns1 = allocator.create_namespace("stats_test1");
    auto ns2 = allocator.create_namespace("stats_test2");

    std::vector<MessageTag> tags;
    for (int i = 0; i < 10; ++i)
    {
        tags.push_back(ns1.request());
        tags.push_back(ns2.request());
    }

    auto final_stats = allocator.get_global_stats();
    EXPECT_EQ(final_stats.total_namespaces, 2);
    EXPECT_EQ(final_stats.total_allocated_tags, 20);
    EXPECT_GT(final_stats.total_available_tags, 0);
    EXPECT_EQ(final_stats.tag_ub_value, allocator.tag_ub());
}

/**
 * 测试命名空间销毁
 */
TEST_F(TagAllocatorTest, NamespaceDestruction)
{
    auto & allocator = Context::for_comm(MPI_COMM_WORLD).tag_allocator();
    auto ns = allocator.create_namespace("destroy_test");
    std::uint32_t ns_id = ns.id();

    // 分配一些标签
    std::vector<MessageTag> tags;
    for (int i = 0; i < 5; ++i)
    {
        tags.push_back(ns.request());
    }

    // 销毁命名空间
    allocator.destroy_namespace(ns_id);

    // 验证标签不再属于任何命名空间
    for (const auto & tag : tags)
    {
        EXPECT_FALSE(allocator.belongs(ns_id, tag.value));
    }

    // 验证无法找到已销毁的命名空间
    EXPECT_THROW(allocator.find_namespace("destroy_test"), std::runtime_error);
}

/**
 * 测试多通信器支持
 */
TEST_F(TagAllocatorTest, MultipleCommSupport)
{
    auto & allocator1 = Context::for_comm(MPI_COMM_WORLD).tag_allocator();
    auto & allocator2 = Context::for_comm(MPI_COMM_WORLD).tag_allocator();

    // 同一通信器应该返回同一实例
    EXPECT_EQ(&allocator1, &allocator2);

    // 不同通信器的测试需要实际的MPI_Comm_dup，这里暂时跳过
    // 因为在单元测试环境中创建新的通信器比较复杂
}